#include "kernel_operator.h"
using namespace AscendC;

#define K_MAX_SHAPE_DIM 0
constexpr int32_t BUFFER_NUM = 2;  
template<typename TYPE_INPUT_X, typename TYPE_OUTPUT_Y> class KernelTrunc {
    public:
        __aicore__ inline KernelTrunc() {}
        __aicore__ inline void Init(GM_ADDR input_x, GM_ADDR output_y,
            uint32_t ALIGN_NUM, uint32_t block_size, uint32_t core_size, uint32_t core_remain,
            TPipe * pipeIn) {
            ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
            this->blockLength = core_size + (GetBlockNum() == GetBlockIdx() + 1 ? core_remain : 0);
            this->tileLength = block_size;
            this->ALIGN_NUM = ALIGN_NUM;
            this->blockLength = this->blockLength + (this->blockLength % ALIGN_NUM ? ALIGN_NUM - this->blockLength % ALIGN_NUM : 0);
            //this->totalLength = totalLength;
            auto startPointer = core_size * GetBlockIdx();
            auto bufferlength = this->blockLength;

    
            Gm_x1.SetGlobalBuffer((__gm__ TYPE_INPUT_X*)input_x + startPointer ,bufferlength);
            Gm_y.SetGlobalBuffer((__gm__ TYPE_OUTPUT_Y*)output_y + startPointer , bufferlength);
            
    
            this->tileNum = this->blockLength / this->tileLength + (this->blockLength % this->tileLength > 0);
    
            pipe = pipeIn;
            pipe->InitBuffer(Q_x1, BUFFER_NUM, this->tileLength * sizeof(TYPE_INPUT_X));

            pipe->InitBuffer(Q_y, BUFFER_NUM, this->tileLength * sizeof(TYPE_INPUT_X));

          
            if constexpr (std::is_same_v<TYPE_INPUT_X, int8_t> || std::is_same_v<TYPE_INPUT_X, uint8_t>) {
               
                pipe->InitBuffer(B_y, this->tileLength * sizeof(half));
            }
            else if constexpr (std::is_same_v<TYPE_INPUT_X, half> ){

                pipe->InitBuffer(B_y, this->tileLength * sizeof(int32_t));
                pipe->InitBuffer(B_y2, this->tileLength * sizeof(float));


                
            }

            else if constexpr (std::is_same_v<TYPE_INPUT_X, bfloat16_t>){

                pipe->InitBuffer(B_y, this->tileLength * sizeof(int32_t));
                pipe->InitBuffer(B_y2, this->tileLength * sizeof(float));


            }

            // (std::is_same_v<TYPE_INPUT_X, float> || std::is_same_v<TYPE_INPUT_X, int32_t> )
            else if constexpr (std::is_same_v<TYPE_INPUT_X, int32_t>){

                pipe->InitBuffer(B_y, this->tileLength * sizeof(float));
             }

            //     pipe.InitBuffer(B_y, this->tileLength * sizeof(float));

            // }
            
        }
        __aicore__ inline void Process() { //假设每次处理一行来看

            int32_t loopCount = this->tileNum;
            for (int32_t i = 0; i < loopCount-1; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            }
            uint32_t length = this->blockLength - this->tileLength * (loopCount - 1);
            CopyIn(loopCount - 1, (length +31 ) / 32 * 32);
            Compute(loopCount - 1, (length +31 ) / 32 * 32);
            CopyOut(loopCount - 1, (length +31 ) / 32 * 32);           
        }   
    
         
    private:
        __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
           LocalTensor<TYPE_INPUT_X> x1 = Q_x1.AllocTensor<TYPE_INPUT_X>();
        //    LocalTensor<TYPE_X1> x2 = Q_x2.AllocTensor<TYPE_X1>();
        //    LocalTensor<uint8_t> cdt = Q_cdt.AllocTensor<uint8_t>();
           DataCopy(x1, Gm_x1[progress * this->tileLength], length);
           //DataCopy(x2, Gm_x2[progress * this->tileLength], length);
           
           //DataCopy(cdt, (Gm_cdt[progress * this->tileLength]), length);
           
           Q_x1.EnQue(x1);
        //    Q_x2.EnQue(x2);
        //    Q_cdt.EnQue(cdt);
        }
        __aicore__ inline void Compute(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_INPUT_X> x1 = Q_x1.DeQue<TYPE_INPUT_X>();
            // LocalTensor<TYPE_X1> x2 = Q_x2.DeQue<TYPE_X1>();
            // LocalTensor<uint8_t> cdt = Q_cdt.DeQue<uint8_t>();
            LocalTensor<TYPE_INPUT_X> y = Q_y.AllocTensor<TYPE_INPUT_X>();

            //auto bits = B_bits.Get<uint8_t>();
            if constexpr (std::is_same_v<TYPE_INPUT_X, half> ) {

                auto p1 = B_y.Get<int32_t>();
                auto p2 = B_y2.Get<float>();

                Cast(p1,x1,RoundMode::CAST_TRUNC, length);//从fp16->int32

                Cast(p2,p1,RoundMode::CAST_TRUNC, length);

                Cast(y,p2,RoundMode::CAST_NONE,length);  //从int32->fp16
            }
            else if constexpr (std::is_same_v<TYPE_INPUT_X, bfloat16_t>){

                auto p1 = B_y.Get<int32_t>();

                auto p2 = B_y2.Get<float>();

                Cast(p1,x1,RoundMode::CAST_TRUNC, length);

                Cast(p2,p1,RoundMode::CAST_TRUNC, length);

                Cast(y,p2,RoundMode::CAST_TRUNC, length);


            }
            else if constexpr ( std::is_same_v<TYPE_INPUT_X, int32_t>) {

                auto p1 = B_y.Get<float>();

                Cast(p1,x1 ,RoundMode::CAST_NONE, length);

                Cast(y,p1,RoundMode::CAST_TRUNC,length);

               
            }
         
            else if constexpr (std::is_same_v<TYPE_INPUT_X, int8_t> || std::is_same_v<TYPE_INPUT_X, uint8_t>)
            {

                auto p1 = B_y.Get<half>();

                Cast(p1,x1 ,RoundMode::CAST_NONE, length);
                Cast(y,p1,RoundMode::CAST_TRUNC,length);

                

            }

            else{

                //std::is_same_v<TYPE_INPUT_X, float> ||

                Cast(y,x1,RoundMode::CAST_TRUNC,length);
            }
            // else{

            //     auto p1 = B_half.Get<half>();

            //     Cast(p1, cdt, RoundMode::CAST_NONE, length);

            //     Compare(bits, p1,zero,CMPMODE::NE,length);

            //     auto float_x1 = B_x1.Get<float>();
            //     auto float_x2 = B_x2.Get<float>();
            //     auto float_y = B_y.Get<float>();
            //     Cast(float_x1, x1, RoundMode::CAST_NONE, length);
            //     Cast(float_x2, x2, RoundMode::CAST_NONE, length);
            //     //Cast(float_y, y, RoundMode::CAST_RINT, length);
            //     Select(float_y, bits, float_x1, float_x2, SELMODE::VSEL_TENSOR_TENSOR_MODE, length);
            //     Cast(y, float_y, RoundMode::CAST_RINT, length);
            // }
            Q_x1.FreeTensor(x1);  
            Q_y.EnQue<TYPE_INPUT_X>(y);
        }
        __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
            LocalTensor<TYPE_INPUT_X> y = Q_y.DeQue<TYPE_INPUT_X>();

           //if constexpr (std::is_same_v<TYPE_INPUT_X, half> || std::is_same_v<TYPE_INPUT_X, bfloat16_t>){

            DataCopy(Gm_y[progress * this->tileLength], y, length);
         
            
            Q_y.FreeTensor(y);
        }
    private:
        TPipe * pipe;
        TQue<QuePosition::VECIN, BUFFER_NUM> Q_x1; 
        TQue<QuePosition::VECOUT, BUFFER_NUM> Q_y;
        TBuf<QuePosition::VECCALC>  B_y,B_y2;

      
        GlobalTensor<TYPE_INPUT_X> Gm_x1, Gm_y;
        uint32_t blockLength;
        uint32_t tileNum;
        //uint32_t totalLength;
        uint32_t tileLength;
        uint32_t ALIGN_NUM;
    };
extern "C" __global__ __aicore__ void trunc(GM_ADDR input_x, GM_ADDR output_y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);

    TPipe pipe;

    KernelTrunc<DTYPE_INPUT_X,DTYPE_OUTPUT_Y> op;
    op.Init(input_x, output_y, tiling_data.ALIGN_NUM, tiling_data.block_size, tiling_data.core_size, tiling_data.core_remain,&pipe);
    op.Process();

    // TODO: user kernel impl
}